from sentence_transformers import SentenceTransformer
from data_indexing import create_index_and_embeddings
from query_module import search_documents
from generation_module import load_generation_pipeline, generate_answer

def rag_query(question, index, documents, embedding_model, generation_pipeline):
    """
    RAG 应用的主入口函数，接收用户问题，检索相关文档，并生成答案。

    Args:
        question: 用户问题 (字符串)
        index: FAISS 向量数据库索引
        documents: 文档列表
        embedding_model: 嵌入模型
        generation_pipeline: QWQ-32B 生成 pipeline

    Returns:
        str: 生成的答案
    """
    relevant_docs = search_documents(question, index, documents, embedding_model) # 检索相关文档
    context = "\n".join(relevant_docs) # 将相关文档拼接成上下文
    answer = generate_answer(question, context, generation_pipeline) # 生成答案
    return answer

if __name__ == '__main__':
    # 模拟企业文档 (实际应用中应从文件或数据库加载)
    documents = [
        "我们的退货政策是客户可以在30天内退货。",
        "2025年新的假期安排，每人享有法定年假12天。",
        "上个月，我们的销售政策是全场商品8折优惠。",
        "公司鼓励员工参与志愿活动，每年提供一天带薪志愿假。"
    ]

    # 加载嵌入模型
    embedding_model = SentenceTransformer('all-MiniLM-L12-v2')

    # 创建向量数据库索引
    index, _ = create_index_and_embeddings(documents)

    # 加载生成 pipeline (注意: 默认是 gpt2 占位符模型，需要替换为 QWQ-32B)
    qg_pipeline = load_generation_pipeline()

    # --- 测试问答 ---
    while True:
        question = input("请输入您的问题 (输入 'exit' 退出): ")
        if question.lower() == 'exit':
            break

        answer = rag_query(question, index, documents, embedding_model, qg_pipeline)
        print(f"\n问题：{question}")
        print(f"答案：{answer}\n")